#!/usr/bin/python2

import urllib2, re, gzip, StringIO, time, copy

Pug_url = "http://pubchem.ncbi.nlm.nih.gov/pug/pug.cgi"
Entrez_url = "http://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi"

Check_tpl = '''<PCT-Data>
    <PCT-Data_input>
        <PCT-InputData>
            <PCT-InputData_request>
                <PCT-Request>
                    <PCT-Request_reqid>%(reqid)s</PCT-Request_reqid>
                    <PCT-Request_type value="status"/>
                </PCT-Request>
            </PCT-InputData_request>
        </PCT-InputData>
    </PCT-Data_input>
</PCT-Data>'''

def gen_download_tpl(opt_list):
    tpl_header = '''<PCT-Data>
    <PCT-Data_input>
        <PCT-InputData>
            <PCT-InputData_download>
                <PCT-Download>
                    <PCT-Download_uids>
                        <PCT-QueryUids>\n'''
    tpl_footer = '''                        </PCT-QueryUids>
                    </PCT-Download_uids>
                    <PCT-Download_format value="sdf"/>
                    <PCT-Download_compression value="gzip"/>
                </PCT-Download>
            </PCT-InputData_download>
        </PCT-InputData>
    </PCT-Data_input>
</PCT-Data>\n'''
    tpl_body = ''
    if 'query_key' in opt_list:
        tpl_body += '''                            <PCT-QueryUids_entrez>
                                <PCT-Entrez>
                                    <PCT-Entrez_db>pccompound</PCT-Entrez_db>
                                    <PCT-Entrez_query-key>%(query_key)s</PCT-Entrez_query-key>
                                    <PCT-Entrez_webenv>%(webenv)s</PCT-Entrez_webenv>
                                </PCT-Entrez>
                            </PCT-QueryUids_entrez>'''
    if 'ids' in opt_list:
        tpl_body += '''                            <PCT-QueryUids_ids>
                                <PCT-ID-List>
                                    <PCT-ID-List_db>pccompound</PCT-ID-List_db>
                                    <PCT-ID-List_uids>
%(ids)s
                                    </PCT-ID-List_uids>
                                </PCT-ID-List>
                            </PCT-QueryUids_ids>\n'''

    return '%s%s%s' %(tpl_header, tpl_body, tpl_footer)

def gen_query_tpl(opt_list):
    assert 'query' in opt_list
    tpl_header = '''<?xml version="1.0"?>
<!DOCTYPE PCT-Data PUBLIC "-//NCBI//NCBI PCTools/EN" "http://pubchem.ncbi.nlm.nih.gov/pug/pug.dtd">
<PCT-Data>
  <PCT-Data_input>
    <PCT-InputData>
      <PCT-InputData_query>
        <PCT-Query>
          <PCT-Query_type>
            <PCT-QueryType>
              <PCT-QueryType_css>
                <PCT-QueryCompoundCS>\n'''
    tpl_footer = '''                  <PCT-QueryCompoundCS_results>%(num_results)s</PCT-QueryCompoundCS_results>
                </PCT-QueryCompoundCS>
              </PCT-QueryType_css>
            </PCT-QueryType>
          </PCT-Query_type>
        </PCT-Query>
      </PCT-InputData_query>
    </PCT-InputData>
  </PCT-Data_input>
</PCT-Data>\n'''

    tpl_body = '''                    <PCT-QueryCompoundCS_query>
                        <PCT-QueryCompoundCS_query_data>%(query)s</PCT-QueryCompoundCS_query_data>
                    </PCT-QueryCompoundCS_query>\n'''

    type_header = '''                 <PCT-QueryCompoundCS_type>\n'''
    type_footer = '''                 </PCT-QueryCompoundCS_type>\n'''
    type_body = ''
    if 'isotope' in opt_list:
        type_body += '''                    <PCT-QueryCompoundCS_type_identical>
                      <PCT-CSIdentity value="same-isotope">%(isotope)s</PCT-CSIdentity>
                    </PCT-QueryCompoundCS_type_identical>\n'''
    if 'similarity' in opt_list:
        type_body += '''                    <PCT-QueryCompoundCS_type_similar>
                        <PCT-CSSimilarity>
                            <PCT-CSSimilarity_threshold>%(similarity)s</PCT-CSSimilarity_threshold>
                        </PCT-CSSimilarity>
                    </PCT-QueryCompoundCS_type_similar>\n'''
    if 'substructure' in opt_list:
        type_body += '''                    <PCT-QueryCompoundCS_type_subss>
                    <PCT-CSStructure>
                        <PCT-CSStructure_bonds value="true"/>
                    </PCT-CSStructure>
                    </PCT-QueryCompoundCS_type_subss>\n'''
    if 'formula' in opt_list:
        type_body += '''                    <PCT-QueryCompoundCS_type_formula>
                    <PCT-CSMolFormula></PCT-CSMolFormula>
                    </PCT-QueryCompoundCS_type_formula>\n'''

    return '%s%s%s%s%s%s' %(tpl_header, tpl_body, type_header,
            type_body, type_footer, tpl_footer)

def sdf2list(sdf_string):
    '''parse the sdf string into list contains info of mols stored in  dicts'''

    l = ""
    line = ""
    sdfs = []
    result_list = []

    for line in sdf_string.replace('\r', '').split('\n'):
        if (line.startswith("$$$$")) and l:
            sdfs.append(l)
            l = ""
        else:
            # some times the first line of the sdf might be $$$$, skip such lines
            if not line.startswith("$$$$"):
                l += line
    if l.replace('\r', '').replace('\n', '').replace('\t', '').replace(' ', '') != '':
        sdfs.append(l)
        l = ""

    for string in sdfs:
        L = []
        tmpl = string.split("> <")
        L = [i.strip().split(">") for i in tmpl]
        L[0].insert(0, "mol_2d_struc")
        try:
            dic = dict(L)
        except:
            L.pop()
            dic = dict(L)
        if dic.has_key('PUBCHEM_COMPOUND_CID'):
            dic['mol_id'] = dic['PUBCHEM_COMPOUND_CID']
        result_list.append(dic)

    return result_list

class Query:
    def __init__(self, pug_url = Pug_url, entrez_url = Entrez_url, debug = False):
        self.pug_url = pug_url
        self.entrez_url = entrez_url
        self.debug = debug
        cookie = urllib2.HTTPCookieProcessor()
        self.opener = urllib2.build_opener(cookie)

    def pug_query(self, val_dic, check_interval = 1):
        arg = ''
        query = ''
        url = self.pug_url
        if val_dic.has_key('term') and val_dic['term'] != '':
            query_dic = {
                    'term': val_dic['term'],
                    'db': 'pccompound',
                    'retmax': val_dic.setdefault('num_results', 10),
                    }
            if val_dic.has_key('query_key') and val_dic.has_key('webenv'):
                query_dic['query_key'] = val_dic['query_key']
                query_dic['WebEnv'] = val_dic['webenv']
                query_dic['WebEnvRq'] = 1
                query_dic['cmd'] = 'Select+from+History'
            url = self.entrez_url
            l = ['%s=%s' %(k,v) for k, v in query_dic.items()]
            arg = '&'.join(l)
            url = ('%s?%s' %(url, arg)).replace(' ', '%20')
            #print query
        else:
            for k in ('formula', 'substructure'):
                if val_dic.has_key(k):
                    val_dic['query'] = val_dic[k]
            if not val_dic.has_key('num_results'):
                val_dic['num_results'] = 10
            query = gen_query_tpl(val_dic.keys()) %val_dic

        if self.debug:
            if query:
                print 'Submited query text: %s' %query
            elif url:
                print 'Url in pug_query: %s' %url

        req = query and urllib2.Request(url, query) or urllib2.Request(url)
        f = self.opener.open(req)
        out = f.read()
        f.close()
        #print out

        while '<PCT-Status value="queued"/>' in out \
                or '<PCT-Status value="running"/>' in out:

            if self.debug:
                print 'PCT-Status: running'

            m = re.search("<PCT-Waiting_reqid>(\d+)</PCT-Waiting_reqid>", out)
            if m:
                reqid = m.group(1)
                out = self.check_query(reqid)
            time.sleep(check_interval)

        return out

    def get_download_url(self, query_dic, check_interval = 1):
        # query_dic, 'query_key' and 'webenv' is needed
        query = gen_download_tpl(query_dic.keys()) %query_dic
        #print query
        req = urllib2.Request(self.pug_url, query)
        f = self.opener.open(req)
        out = f.read()
        f.close()
        while '<PCT-Status value="queued"/>' in out or '<PCT-Status value="running"/>' in out:
            m = re.search( "<PCT-Waiting_reqid>(\d+)</PCT-Waiting_reqid>", out)
            if m:
                reqid = m.group(1)
                out = self.check_query(reqid)
            time.sleep(check_interval)

        url = re.search("<PCT-Download-URL_url>(.*)</PCT-Download-URL_url>", out).group(1)
        return url

    def check_query(self, reqid):
        # check the query state, returns the report text
        check = Check_tpl %{'reqid': reqid}
        req = urllib2.Request(self.pug_url, check)
        f = self.opener.open(req)
        out = f.read()
        f.close()
        return out

    def get_return_info(self, text):

        result = {}
        if 'IdList' in text:
            l = re.findall('<Id>(\d+)</Id>', text)
            for id in l:
                if not result.has_key('ids'):
                    result['ids'] = ''
                result['ids'] += '''                                        <PCT-ID-List_uids_E>%s</PCT-ID-List_uids_E>\n''' %id
        elif not 'eSearchResult' in text:
            try:
                m = re.search("<PCT-Entrez_query-key>(\d+)</PCT-Entrez_query-key>", text)
                result['query_key'] = m.group(1)
                m = re.search("<PCT-Entrez_webenv>(.+?)</PCT-Entrez_webenv>", text)
                result['webenv'] = m.group(1)
            except:
                pass
        return result

    def download(self, url):
        stream = ''
        try:
            req = urllib2.Request(url)
            sdf = self.opener.open(req)
            zbuf = StringIO.StringIO(sdf.read())
            stream = gzip.GzipFile(fileobj = zbuf, mode = 'rb').read()
            sdf.close()
        except:
            pass
        return stream

    def search(self, val_dic, download = False, check_interval = 1):
        # val_dic: 'query' or 'term' key is needed.
        # For 'query', searching using SMILES and formula is supported.
        # The 'query' key also accept valid CID number as its value
        # Similarity search and substructure search can also be performed
        # use 'isotope' key to perform exact search.
        # For 'term', see http://pubchem.ncbi.nlm.nih.gov/search/help_search.html#TextSch
        # and valid terms can be found at: http://pubchem.ncbi.nlm.nih.gov/help.html#idxcompound
        # Examples:
        #   Exact search:
        #       val_dic = {'query': 'c1ccccc1C', 'num_results': 100, 'isotope': 4}
        #   Similarity search:
        #       val_dic = {'query': 'c1ccccc1C', 'num_results': 100, 'similarity': 80}
        #       80 is the Tanimoto similarity value in percentage
        #   Formula search:
        #       val_dic = {'query': 'omittable', 'num_results': 100, 'formula': 'CH3COOH'}
        #   Substructure search:
        #       val_dic = {'query': 'omittable', 'num_results': 100, 'substructure': 'c1ccccc1C'}
        result_dic = {}
        combined_search = False
        query_dic = copy.deepcopy(val_dic)
        new_query_dic = copy.deepcopy(val_dic)
        returned_query_dic = {}
        #open('query.txt', 'w').write(str(val_dic))
        if 'term' in val_dic.keys():
            for k in ('substructure', 'formula', 'isotope', 'similarity'):
                if k in val_dic.keys():
                    if not combined_search:
                        combined_search = True
                        query_dic.pop('term')
                    new_query_dic.pop(k)

        result = self.pug_query(query_dic, check_interval)
        #print result
        returned_query_dic = self.get_return_info(result)

        if combined_search:
            new_query_dic.update({'query_key': returned_query_dic['query_key'],
                'webenv': returned_query_dic['webenv']})
            result = self.pug_query(new_query_dic, check_interval)
            if self.debug:
                print 'Query key: %s \nWeb env: %s' %(returned_query_dic['query_key'],
                        returned_query_dic['webenv'])
            returned_query_dic = self.get_return_info(result)

        #print query_dic
        if returned_query_dic:
            url = self.get_download_url(returned_query_dic)
        else:
            result_dic['id'] = val_dic.setdefault('id', '')
            result_dic['mol'] = val_dic.setdefault('mol', '')
            result_dic['info'] = ''
            result_dic['len_results'] = 0
            result_dic['result_list'] = []
            result_dic['sdf'] = ''
            return result_dic

        if self.debug:
            print 'Query dict: %s' %str(returned_query_dic)
            print 'Result URL: %s' %url

        if download:
            sdf = self.download(url)
            tmpl = sdf2list(sdf)
            result_dic['id'] = val_dic.setdefault('id', '')
            result_dic['mol'] = val_dic.setdefault('mol', '')
            result_dic['info'] = ''
            result_dic['len_results'] = len(tmpl)
            result_dic['result_list'] = tmpl
            result_dic['sdf'] = sdf
            return result_dic
        return True

if __name__ == '__main__':
    # Examples:
    #   Exact search:
    #       val_dic = {'query': 'c1ccccc1C', 'num_results': 100, 'isotope': 4}
    #   Similarity search:
    #       val_dic = {'query': 'c1ccccc1C', 'num_results': 100, 'similarity': 80}
    #       80 is the Tanimoto similarity value in percentage
    #   Formula search:
    #       val_dic = {'query': 'omittable', 'num_results': 100, 'formula': 'CH3COOH'}
    #   Substructure search:
    #       val_dic = {'query': 'omittable', 'num_results': 100, 'substructure': 'c1ccccc1C'}

    query = Query()

    val_dic = {'substructure': 'c1ccccc1C', 'term': '500:600[MW]', 'num_results': 200}
    #val_dic = {'term': '100:100.01[exactmass]', 'num_results': 100}
    #val_dic = {'num_results': 100, 'formula': 'CH3COOH'}
    #val_dic = {'query': 'C([C@@H]1C[C@@H](CCC)CN1C)(N[C@@H](C1[C@@H]([C@@H]([C@H]([C@H](O1)SC)O)O)O)[C@@H](Cl)C)=O', 'isotope': '4'}
    #val_dic = {'query': '[C@@H]1(N(C)C[C@H](CCC)C1)C(=O)N[C@H]([C@@H](Cl)C)C1O[C@H](SC)[C@@H]([C@H]([C@H]1O)O)O', 'isotope': 4}
    #val_dic = {'term': '602:602.3754[EXMASS] AND 0[CHG]', 'num_results': 10000, 'substructure': 'C1(O)=CC(=O)C2(C=O)C(C)(C)CCC1C2=O', 'id': '77529-1'}

    #val_dic = {'query': 'C[C@H](Cl)[C@@H](NC([C@H]1N(C)C[C@H](CCC)C1)=O)C1O[C@H](SC)[C@H](O)[C@H]([C@H]1O)O', 'num_results': 10000, 'id': '19874-1', 'isotope': 4}

    val_dic = {'term': '601.907825:602.907825[EXMASS] AND 0[CHG]', 'num_results': 10000, 'substructure': 'C(=O)C1C(=O)C2CCC(C)C(C1=O)C2=O', 'id': '86947-1'}

        # turn on debug mode
    query.debug = True
    #print query.get_return_info(s)
    #url='http://eutils.ncbi.nlm.nih.gov/entrez/eutils/esearch.fcgi?query_key=1&term=602:602.3754[EXMASS]%20AND%200[CHG]&retmax=10000&WebEnv=NCID_1_140702096_130.14.22.28_9001_1329818605_990547547&WebEnvRq=1&cmd=Select+from+History&db=pccompound'
    #cookie = urllib2.HTTPCookieProcessor()
    #opener = urllib2.build_opener(cookie)
    #req = urllib2.Request(url)
    #f = opener.open(req)
    #print f.read()
    #f.close()


    print query.search(val_dic, True)
    #print sdf2list(open('temp.sdf', 'r').read())

